perf(bagel): batch trainside UniGRPO rollout forwards - #253

Summary

• add opt-in block-diagonal batching for BAGEL AR thinking chains and same-shape diffusion samples
• adapt packed diffusion to the current per-sample KV-in-conditions representation
• keep forward_batch_size=1 behavior and RNG semantics unchanged by default
• preserve serial fallbacks for ViT AR inputs, deferred contexts, and mixed image shapes

Motivation

BAGEL unified full-FT spends about 76% of a G=4 probe step in generation; at the production G=24 fanout this is expected to approach 90–95%. The current trainside path executes each sibling with navit batch size one and repeats live-FSDP all-gathers.

Fresh 8×H20 full-FT A/B

Identical config for all arms: batch_size=8, samples_per_prompt=4, max_new_tokens=512, two rollouts, no checkpoint save. Steady time is the rollout-1 → rollout-2 timestamp delta.

Pack-4 captures almost all of the gain and is the safer memory-tuning recommendation; the committed default remains 1.

Test plan

• focused fake-model parity/isolation/RNG tests (7 passed)
• Ruff lint/format and repository pre-push hooks
• serial forward_batch_size=1 keeps the legacy no-generator path
• fresh baseline vs pack-4 vs pack-8 full-FT A/B on 8×H20
• all arms completed two optimizer steps without OOM
• image rollout/replay ratio remained 1.0000±0.0000

Historical prototype reference: 2×8 steady step improved from ~267s to ~132s; this PR is a clean port onto current main/KV conditions rather than a rebase of that stale branch.

Maintainer validation (347a22e)

The seven focused test_pack_b cases passed before the PR-added test file was removed per repository preference. No tests/** addition remains in the PR. compileall, full pre-commit, and recipe target resolution pass on the current head. Runtime risk remains moderate/high despite the opt-in default because the packed BAGEL path changes model execution and has only the author hardware A/B evidence.